import torch
import torch.nn as nn
import torch.nn.functional as F

def do_make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
    """Make mask tensor containing indices of padded part.

    See description of make_non_pad_mask.

    Args:
        lengths (torch.Tensor): Batch of lengths (B,).
    Returns:
        torch.Tensor: Mask tensor containing indices of padded part.

    Examples:
        >>> lengths = [5, 3, 2]
        >>> make_pad_mask(lengths)
        masks = [[0, 0, 0, 0 ,0],
                 [0, 0, 0, 1, 1],
                 [0, 0, 1, 1, 1]]
    """
    batch_size = lengths.size(0)
    max_len = max_len if max_len > 0 else lengths.max().item()
    seq_range = torch.arange(0,
                             max_len,
                             dtype=torch.int64,
                             device=lengths.device)
    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
    seq_length_expand = lengths.unsqueeze(-1)
    mask = seq_range_expand >= seq_length_expand
    return mask

def do_make_seq_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
    """Make mask tensor containing indices of padded part.

    See description of make_non_pad_mask.

    Args:
        lengths (torch.Tensor): Batch of lengths (B,).
    Returns:
        torch.Tensor: Mask tensor containing indices of padded part.

    Examples:
        >>> lengths = [5, 3, 2]
        >>> make_pad_mask(lengths)
        masks = [[1, 1, 1, 1 ,1],
                 [1, 1, 1, 0, 0],
                 [1, 1, 0, 0, 0]]
                 ]
    """
    return ~do_make_pad_mask(lengths, max_len)

class DeterministicDownsampleSplice(nn.Module):
    """
    确定性拼接下采样模块：将连续N帧特征拼接为1帧，实现时序下采样并扩大特征维度
    适用场景：语音、时序信号等需要保留多帧上下文的任务
    """
    def __init__(self, group_size=3):
        """
        Args:
            group_size: 下采样因子（每group_size帧拼接为1帧）
        """
        super().__init__()
        self.group_size = group_size  # 每组帧数（下采样率）

    def forward(self, x):
        """
        前向传播：对输入时序特征进行拼接下采样
        Args:
            x: 输入张量，形状为 (Batch, Seq, Dim)
                - Batch：批次大小
                - Seq：时序长度（帧数）
                - Dim：每帧特征维度
        Returns:
            下采样后的张量，形状为 (Batch, Seq//group_size, group_size*Dim)
        """
        # 解析输入形状
        B, T, D = x.size()  # B=批次, T=时序长度, D=特征维度

        # 计算填充长度：确保时序长度能被group_size整除（避免无法均分）
        pad_len = (self.group_size - T % self.group_size) % self.group_size
        if pad_len > 0:
            # 在时序维度（第2维）右侧填充，使用replicate模式（复制最后一帧的值）
            # F.pad参数：(左, 右, 上, 下)，这里仅填充时序维度的右侧（右=pad_len）
            x = F.pad(x.float(), (0, 0, 0, pad_len), mode='replicate').to(x.dtype)
            T += pad_len  # 更新填充后的时序长度
        # --------------------------
        # 4. 特征下采样（拼接操作）
        # --------------------------
        return x.contiguous().view(B, T // self.group_size, self.group_size * D)

def downsample_seq_mask(seq_mask, group_size):
    """
    对序列掩码进行下采样（与拼接下采样同步）
    Args:
        seq_mask: 原始序列掩码，形状 (Batch, Seq)，1=有效，0=填充
        group_size: 下采样因子（每group_size帧拼接为1帧）
    Returns:
        down_mask: 下采样后的序列掩码，形状 (Batch, Down_Seq)
    """
    B, T = seq_mask.shape  # B=批次，T=原始时序长度
    device = seq_mask.device  # 保持设备一致性

    # 1. 计算需要填充的长度（确保原始长度能被group_size整除）
    pad_len = (group_size - T % group_size) % group_size
    if pad_len > 0:
        # 在时序维度右侧填充0（填充部分视为无效）
        seq_mask = F.pad(seq_mask, (0, pad_len), mode='constant', value=0)
        T += pad_len  # 更新填充后的长度

    # 2. 按group_size分组，每组内取最大值（只要有1个有效，结果就为1）
    down_seq = T // group_size  # 下采样后的时序长度
    # 重塑为 (B, down_seq, group_size)，然后在每组内取max
    down_mask = seq_mask.view(B, down_seq, group_size).max(dim=2)[0]

    return down_mask.to(device)

class SpeechProjectionModule(nn.Module):
    def __init__(self, speech_embedding_dim, lm_hidden, downsample_rate=2):
        super().__init__()
        self.speech_embedding_dim = speech_embedding_dim  # 原始语音特征维度
        self.lm_hidden = lm_hidden  # 投影后的目标维度
        self.downsample_rate = downsample_rate  # 下采样因子

        # 构建下采样+投影的序列模块
        self.speech_projection = nn.Sequential(
            # 第一步：拼接下采样（输入x和x_length，输出下采样特征、长度、掩码）
            DeterministicDownsampleSplice(group_size=downsample_rate),
            # 第二步：线性投影（将拼接后的维度映射到目标维度）
            nn.Linear(
                speech_embedding_dim * downsample_rate,  # 拼接后的维度
                lm_hidden  # 目标维度
            )
        )

    def forward(self, x, seq_mask):
        """
        Args:
            x: 语音特征，形状 (Batch, Seq, speech_embedding_dim)
            x_length: 语音序列长度，形状 (Batch,)
        Returns:
            projected_x: 投影后的特征，形状 (Batch, Down_Seq, lm_hidden)
            down_length: 下采样后的序列长度，形状 (Batch,)
            seq_mask: 下采样后的序列掩码，形状 (Batch, Down_Seq)
        """
        x = self.speech_projection(x)
        seq_mask_new = downsample_seq_mask(seq_mask, self.downsample_rate)
        return x, seq_mask_new

# --------------------------
# 使用示例：验证拼接下采样的效果
# --------------------------
if __name__ == "__main__":
    # 1. 设置参数
    group_size = 3  # 下采样因子：每3帧拼接为1帧
    batch_size = 2  # 批次大小
    seq_len = 10    # 原始时序长度（帧数）
    feat_dim = 8   # 每帧特征维度

    # 2. 生成随机测试数据（模拟时序特征，如语音的梅尔频谱）
    # 形状：(Batch, Seq, Dim) = (2, 10, 80)
    x = torch.randn(batch_size, seq_len, feat_dim)
    x_lens = torch.tensor([10, 9])  # 原始序列长度
    seq_mask = do_make_seq_mask(x_lens)
    print(f"原始输入形状：{x.shape}")
    print(f"下采样因子（group_size）：{group_size}")
    print(f"原始序列长度：{x_lens}")
    print(f"原始序列掩码：{seq_mask}")
    # 3. 初始化并应用下采样模块
    downsampler = SpeechProjectionModule(speech_embedding_dim=feat_dim, lm_hidden=feat_dim, downsample_rate=group_size)
    output, seq_mask = downsampler(x, seq_mask)

    # 4. 输出结果分析
    print(f"下采样后输出形状：{output.shape}")
    print(f"下采样后序列掩码：{seq_mask}")
    print("\n形状变化解析：")
    print(f"- 批次大小（Batch）：保持不变 → {batch_size}")
    print(f"- 时序长度（Seq）：{seq_len} → {output.shape[1]} "
          f"（因原始长度10不能被3整除，填充2帧后为12，12//3=4）")
    print(f"- 特征维度（Dim）：{feat_dim} → {output.shape[2]} "
          f"（扩大为group_size倍：3×80=240）")

    # 5. 验证拼接逻辑（取第一批次的前几帧手动对比）
    print("\n拼接逻辑验证：")
    print("原始第一批次前3帧的前2个特征值：")
    print(x[0, 0, :])  # 第1帧前2维
    print(x[0, 1, :])  # 第2帧前2维
    print(x[0, 2, :])  # 第3帧前2维
    print("下采样后第一批次第1帧的前6个特征值（应为上述3帧的拼接）：")
    print(output[0, 0, :])  # 拼接后前6维 = 第1帧前2维 + 第2帧前2维 + 第3帧前2维
    print(output[0, -1, :])  # 拼接后前6维 = 第1帧前2维 + 第2帧前2维 + 第3帧前2维